#include <string.h>
#include "sm3.h"

#define T0 0x79cc4519U
#define T1 0x7a879d8aU

#define rlft(x, n)      (((x) << ((n) % 32)) | ((x) >> (32 - (n))))
#define P0(x)           ((x) ^ rlft((x),  9) ^ rlft((x), 17))
#define P1(x)           ((x) ^ rlft((x), 15) ^ rlft((x), 23))
#define FF0(x, y, z)    (x ^ y ^ z)
#define FF1(x, y, z)    ((x & y) | (x & z) | (y & z))
#define GG0(x, y, z)    (x ^ y ^ z)
#define GG1(x, y, z)    (((y ^ z) & x) ^ z)

#define RF0(a, b, c, d, e, f, g, h, i)                  \
        ss1 = rlft(a, 12) + e + rlft(T0, i % 32);       \
        ss1 = rlft(ss1, 7);                             \
        ss2 = rlft(a, 12) ^ ss1;                        \
        d += FF0(a, b, c) + ss2 + (w[i] ^ w[i + 4]);    \
        h += GG0(e, f, g) + ss1 + w[i];                 \
        h = P0(h);                                      \
        b = rlft(b, 9);                                 \
        f = rlft(f, 19);

#define RF1(a, b, c, d, e, f, g, h, i)                  \
        ss1 = rlft(a, 12) + e + rlft(T1, i % 32);       \
        ss1 = rlft(ss1, 7);                             \
        ss2 = rlft(a, 12) ^ ss1;                        \
        d += FF1(a, b, c) + ss2 + (w[i] ^ w[i + 4]);    \
        h += GG1(e, f, g) + ss1 + w[i];                 \
        h = P0(h);                                      \
        b = rlft(b, 9);                                 \
        f = rlft(f, 19);


void sm3_init(u32 buf[8]) {
    buf[0] = 0x7380166f;
    buf[1] = 0x4914b2b9;
    buf[2] = 0x172442d7;
    buf[3] = 0xda8a0600;
    buf[4] = 0xa96f30bc;
    buf[5] = 0x163138aa;
    buf[6] = 0xe38dee4d;
    buf[7] = 0xb0fb0e4e;
}

void sm3_compress(const u8 *msqblock, u32 buf[8]) {
    register u32 a = buf[0], b = buf[1], c = buf[2], d = buf[3], e = buf[4], f = buf[5],
    g = buf[6], h = buf[7], ss1, ss2;
    u32 w[68];

    //消息扩展
    #pragma unroll
    for (int i = 0; i < 16; i++) {
        w[i] = Rev(msqblock + (i << 2));
    }
    #pragma unroll
    for (int i = 16; i < 68; i++) {
        w[i] = P1(w[i - 16] ^ w[i - 9] ^ rlft(w[i - 3], 15)) ^ rlft(w[i - 13], 7) ^ w[i - 6];
    }
    RF0(a, b, c, d, e, f, g, h, 0)
    RF0(d, a, b, c, h, e, f, g, 1)
    RF0(c, d, a, b, g, h, e, f, 2)
    RF0(b, c, d, a, f, g, h, e, 3)
    RF0(a, b, c, d, e, f, g, h, 4)
    RF0(d, a, b, c, h, e, f, g, 5)
    RF0(c, d, a, b, g, h, e, f, 6)
    RF0(b, c, d, a, f, g, h, e, 7)
    RF0(a, b, c, d, e, f, g, h, 8)
    RF0(d, a, b, c, h, e, f, g, 9)
    RF0(c, d, a, b, g, h, e, f, 10)
    RF0(b, c, d, a, f, g, h, e, 11)
    RF0(a, b, c, d, e, f, g, h, 12)
    RF0(d, a, b, c, h, e, f, g, 13)
    RF0(c, d, a, b, g, h, e, f, 14)
    RF0(b, c, d, a, f, g, h, e, 15)
    RF1(a, b, c, d, e, f, g, h, 16)
    RF1(d, a, b, c, h, e, f, g, 17)
    RF1(c, d, a, b, g, h, e, f, 18)
    RF1(b, c, d, a, f, g, h, e, 19)
    RF1(a, b, c, d, e, f, g, h, 20)
    RF1(d, a, b, c, h, e, f, g, 21)
    RF1(c, d, a, b, g, h, e, f, 22)
    RF1(b, c, d, a, f, g, h, e, 23)
    RF1(a, b, c, d, e, f, g, h, 24)
    RF1(d, a, b, c, h, e, f, g, 25)
    RF1(c, d, a, b, g, h, e, f, 26)
    RF1(b, c, d, a, f, g, h, e, 27)
    RF1(a, b, c, d, e, f, g, h, 28)
    RF1(d, a, b, c, h, e, f, g, 29)
    RF1(c, d, a, b, g, h, e, f, 30)
    RF1(b, c, d, a, f, g, h, e, 31)
    RF1(a, b, c, d, e, f, g, h, 32)
    RF1(d, a, b, c, h, e, f, g, 33)
    RF1(c, d, a, b, g, h, e, f, 34)
    RF1(b, c, d, a, f, g, h, e, 35)
    RF1(a, b, c, d, e, f, g, h, 36)
    RF1(d, a, b, c, h, e, f, g, 37)
    RF1(c, d, a, b, g, h, e, f, 38)
    RF1(b, c, d, a, f, g, h, e, 39)
    RF1(a, b, c, d, e, f, g, h, 40)
    RF1(d, a, b, c, h, e, f, g, 41)
    RF1(c, d, a, b, g, h, e, f, 42)
    RF1(b, c, d, a, f, g, h, e, 43)
    RF1(a, b, c, d, e, f, g, h, 44)
    RF1(d, a, b, c, h, e, f, g, 45)
    RF1(c, d, a, b, g, h, e, f, 46)
    RF1(b, c, d, a, f, g, h, e, 47)
    RF1(a, b, c, d, e, f, g, h, 48)
    RF1(d, a, b, c, h, e, f, g, 49)
    RF1(c, d, a, b, g, h, e, f, 50)
    RF1(b, c, d, a, f, g, h, e, 51)
    RF1(a, b, c, d, e, f, g, h, 52)
    RF1(d, a, b, c, h, e, f, g, 53)
    RF1(c, d, a, b, g, h, e, f, 54)
    RF1(b, c, d, a, f, g, h, e, 55)
    RF1(a, b, c, d, e, f, g, h, 56)
    RF1(d, a, b, c, h, e, f, g, 57)
    RF1(c, d, a, b, g, h, e, f, 58)
    RF1(b, c, d, a, f, g, h, e, 59)
    RF1(a, b, c, d, e, f, g, h, 60)
    RF1(d, a, b, c, h, e, f, g, 61)
    RF1(c, d, a, b, g, h, e, f, 62)
    RF1(b, c, d, a, f, g, h, e, 63)
    buf[0] ^= a;
    buf[1] ^= b;
    buf[2] ^= c;
    buf[3] ^= d;
    buf[4] ^= e;
    buf[5] ^= f;
    buf[6] ^= g;
    buf[7] ^= h;
}

void sm3(u8 digest[32], u8 *msg, u64 msglen) {
    u32 buf[8] = {0};
    u64 nblock = msglen >> 9;
    sm3_init(buf);
    for (u32 i = 0; i < nblock; i++)
        sm3_compress(msg + (i << 6), buf);

    // 填充1000**
    u8 tmp[64] = {0};
    u32 rbit= msglen & 511;
    if (0 < rbit && rbit < 448) {                       //最后一个消息块
        memcpy(tmp, msg + (nblock << 6), rbit >> 3);
        tmp[rbit >> 3] = (rbit & 7) ? (msg[msglen >> 3] | (0x80 >> (msglen & 7))) : 0x80;
    } else if (rbit >= 448) {                           //最后两个消息块，填充的1在第一块
        memcpy(tmp, msg + (nblock << 6), rbit >> 3);
        tmp[rbit >> 3] = (rbit & 7) ? (msg[msglen >> 3] | (0x80 >> (msglen & 7))) : 0x80;
        sm3_compress(tmp, buf);
        memset(tmp, 0, 56);
    } else                                              //最后两个消息块，填充的1在第二块
        tmp[0] = 0x80;


    //长度填充
    u8 *len_b = (u8 *) &msglen;
    for (int i = 0; i < 8; i++)
        tmp[56 + i] = len_b[7 - i];
    sm3_compress(tmp, buf);

    //输出
    Put(digest + 0, buf[0]);
    Put(digest + 4, buf[1]);
    Put(digest + 8, buf[2]);
    Put(digest + 12, buf[3]);
    Put(digest + 16, buf[4]);
    Put(digest + 20, buf[5]);
    Put(digest + 24, buf[6]);
    Put(digest + 28, buf[7]);
    memset(buf, 0, 32);
}

